import logging
from datetime import datetime
from tqdm import tqdm
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from pinecone import Pinecone, ServerlessSpec
import pinecone
import numpy as np
from collections import Counter
from sklearn.datasets import load_digits
import matplotlib.pyplot as plt

# 配置日志记录器
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

pinecone = Pinecone(api_key="acd1bdc5-a876-42db-8b49-191bd4bb1065")
index_name = "mnist-index"

# 获取现有索引列表
existing_indexes = pinecone.list_indexes()

# 检查索引是否存在，如果存在就删除
if any(index['name'] == index_name for index in existing_indexes):
    logging.info(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已成功删除。")
else:
    logging.info(f"索引 '{index_name}' 不存在，将创建新索引。")

# 加载MNIST数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建新索引
logging.info(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
logging.info(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
logging.info(f"已成功连接到索引 '{index_name}'。")

# 初始化一个空列表，用于存储转换后的向量数据
vectors = []

# 遍历所有样本，将数据转换为 Pinecone 可接受的格式
for i in range(len(X_train)):
    vector_id = str(i)
    vector_values = X_train[i].tolist()
    metadata = {"label": int(y_train[i])}
    vectors.append((vector_id, vector_values, metadata))

# 定义批处理大小，每批最多包含 1000 个向量
batch_size = 1000

# 使用步长为 batch_size 的 range 函数，实现分批处理
for i in range(0, len(vectors), batch_size):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)
    logging.info(f"已上传 {len(batch)} 条数据到索引 '{index_name}'。")

# 使用测试集进行预测并计算准确率
correct_predictions = 0
total_predictions = len(X_test)

for i in tqdm(range(len(X_test)), desc="Testing", unit="sample"):
    query_data = X_test[i].tolist()
    results = index.query(
        vector=query_data,
        top_k=11,
        include_metadata=True
    )
    labels = [match['metadata']['label'] for match in results['matches']]
    final_prediction = Counter(labels).most_common(1)[0][0]
    if final_prediction == y_test[i]:
        correct_predictions += 1

accuracy = correct_predictions / total_predictions * 100
logging.info(f"当 k=11 时，Pinecone 的准确率为：{accuracy:.2f}%")
